import cv2
import numpy as np
import onnxruntime
# from ocr.utils.try_exception import error_handler1


class DocumentImagesClassifier:
    def __init__(self, model_path):
        self.input_size = (224, 224)
        self.means = (0.485, 0.456, 0.406)
        self.stds = (0.229, 0.224, 0.225)

        self.label_dicts = {'0': '发票', '1': '套打手写', '2': '出生证明', '3': '户口本',
                            '4': '银行卡', '5': '身份证正面', '6': '身份证背面', '7': '其他'}

        self.sess = onnxruntime.InferenceSession(model_path)
        self.input_name = self.sess.get_inputs()[0].name

    def resize_image(self, ori_image):
        ori_size = ori_image.shape[0:2]

        ratio = min(float(self.input_size[i]) / (ori_size[i]) for i in range(len(ori_size)))
        new_size = tuple([int(i * ratio) for i in ori_size])

        img_resized = cv2.resize(ori_image, (new_size[1], new_size[0]), interpolation=3)
        pad_w = self.input_size[1] - new_size[1]
        pad_h = self.input_size[0] - new_size[0]
        top, bottom = pad_h // 2, pad_h - (pad_h // 2)
        left, right = pad_w // 2, pad_w - (pad_w // 2)
        img_resized = cv2.copyMakeBorder(img_resized, top, bottom, left, right, cv2.BORDER_CONSTANT,
                                         value=(104, 116, 124))

        return img_resized

    def pre_process_image(self, img):
        img_resized = self.resize_image(img)
        img_resized = img_resized[:, :, [2, 1, 0]]

        img = img_resized.astype(np.float32) / 255
        for i in range(len(self.means)):
            img[:, :, i] = (img[:, :, i] - self.means[i]) / self.stds[i]

        img = np.array([img])

        return img.transpose((0, 3, 1, 2))

    def post_process_result(self, result):
        label = self.label_dicts[str(result[0].argmax(axis=1)[0])]

        return label

    def predict(self, img):
        input_img = self.pre_process_image(img)

        res = self.sess.run(None, {self.input_name: input_img})

        label = self.post_process_result(res)

        return [{"image_class": label}]
